#pragma once

#include <vector>
#include "common\define.h"
#include "NvInfer.h"
#include "cuda_img_resize.cuh"

namespace yolo_v6
{
	struct DetectRes2 {
		int classes;
		float x;
		float y;
		float w;
		float h;
		float prob;
	};

	class YoloDet2
	{
	public:
		

	public:
		YoloDet2();
		~YoloDet2();
	
		bool Init(int gpu_id, const char* hdats_path);
		void Detect(st_dlm_data* p_imgs, int num, std::vector<std::vector<DetectRes2>> &vv_results);
		int BatchSize();
	private:
		void Reset();
		bool LoadEngine(std::vector<char> v_engine_data);
		std::vector<std::vector<DetectRes2>> postProcess(st_dlm_data* p_imgs,  int num, float* output, const int& outSize);
		void NmsDetect(std::vector <DetectRes2>& detections, float nms_threshold);
		float IOUCalculate(const DetectRes2& det_a, const DetectRes2& det_b);
		std::vector<float> prepareImage(st_dlm_data* p_imgs, int num);

		st_encrypt_info m_info;
		nvinfer1::ICudaEngine* m_engine = nullptr;
		nvinfer1::IExecutionContext* m_context = nullptr;
		std::vector<void*> m_v_dev_buffer;
		std::vector<int64_t> m_v_dev_buf_size;

		std::vector<std::vector<int>> m_grids;

		int m_out_size = 0;
		cudaStream_t m_stream;
		std::vector<float> m_v_out_buf;
		std::vector<st_cuda_resize_dev_space> _v_resize_space;
	};

}//namespace yolo_v6